#!/usr/bin/env python
from __future__ import print_function
import argparse
import pysam
import numpy as np
import pandas as pd


def get_args():
    ap = argparse.ArgumentParser(description='description')
    ap.add_argument('bam', help='input bam file')
    ap.add_argument('group', help='group file')
    ap.add_argument('out', help='output bam file')
    # ap.add_argument('sub', nargs='?', help='sub')
    ap.add_argument('-p', '--prefix', help='out prefix', default='1')
    ap.add_argument('-c', '--coverage', help='min coverage', default=2, type=int)
    # ap.add_argument('-r', '--rename', help='False default', action="store_true")
    return ap.parse_args()


def get_tag(read, tag):
    try:
        return read.get_tag(tag)
    except KeyError as e:
        return None


def filter_ug_count(reads, ug='UG', min_cov=2):
    ug = [get_tag(i, ug) for i in reads]
    ug = pd.Series(ug)
    ug_count = ug.value_counts()

    # ug_count_freq = ug_count.value_counts()
    ug_count_filt = ug_count[ug_count >= min_cov]
    ug_filt = ug.apply(lambda x: x in ug_count_filt.index)

    reads = [reads[i] for i in ug_filt.index]


def is_paired(read1, read2):
    if read1.query_name == read2.query_name:
        return True
    else:
        return False


def tes(bam, group, out, min_cov=2):
    samfile = pysam.AlignmentFile(bam, 'rb')
    bam_filt = pysam.AlignmentFile(out, 'wb', template=samfile)

    reads = []
    for read in samfile.fetch():
        reads.append(read)

    df_group = get_group(group)
    ids = filter_read_group(df_group, min_cov)
    reads_filt = filter_reads_by_id(reads, ids)

    for r in reads_filt:
        bam_filt.write(r)


def get_group(path):
    return pd.read_table(path, index_col=0)


def filter_read_group(df, min_cov=2):
    print(df.head(2))
    final_umi_count = 'final_umi_count'
    read_id = 'read_id'
    df = df.loc[df[final_umi_count] >= min_cov, :]
    return df.index


def filter_reads_by_id(reads, ids):
    df = pd.DataFrame(reads, columns=['read'])
    df.index = [i.query_name for i in reads]
    print(ids[:2])
    print(df.head(2))
    ids = ids.intersection(df.index).unique()
    df = df.loc[ids, :]
    return df.iloc[:, 0].tolist()


def main():
    args = get_args()
    tes(args.bam, args.group, args.out, args.coverage)


if __name__ == '__main__':
    main()
